#include "ir.h"
#include <cassert>
#include <iostream>
using namespace ir;
using namespace std;

static string OperandToString(Operand& operand)
{
    if (holds_alternative<int>(operand)) 
    { return to_string(get<int>(operand)); }
    else { return std::get<std::string>(operand); }
}

string Instruction::ToString()
{
    string eq_str = (dest == "")? "" : " = ";
    string op_str;
    switch (op)
    {
        case OP_ADD: op_str = "add"; break;
        case OP_SUB: op_str = "sub"; break;
        case OP_MUL: op_str = "mul"; break;
        case OP_DIV: op_str = "div"; break;
        case OP_MOD: op_str = "mod"; break;
        case OP_NE: op_str = "ne"; break;
        case OP_EQ: op_str = "eq"; break;
        case OP_GT: op_str = "gt"; break;
        case OP_LT: op_str = "lt"; break;
        case OP_GE: op_str = "ge"; break;
        case OP_LE: op_str = "le"; break;
        case OP_AND: op_str = "and"; break;
        case OP_OR: op_str = "or"; break;
        case OP_XOR: op_str = "xor"; break;
        case OP_SHL: op_str = "shl"; break;
        case OP_SHR: op_str = "shr"; break;
        case OP_SAR: op_str = "sar"; break;
        case OP_LOAD: op_str = "load"; break;
        case OP_STORE: op_str = "store"; break;
        case OP_ALLOC: op_str = "alloc i32"; break;
        case OP_MV: op_str = "mv"; break;
        case OP_BR: op_str = "br"; break;
        case OP_JMP: op_str = "jmp"; break;
        case OP_RET: op_str = "ret"; break;
        case OP_CALL: op_str = "call"; break;
        default: assert(false); break;
    } 
    string srcs_str = "";
    size_t i = 0;
    if (op == OP_CALL)
    {
        assert(srcs.size() > 0);
        srcs_str += OperandToString(srcs[0]) + "(";
        ++i;
    }
    for (; i < srcs.size(); ++i)
    {
        srcs_str += OperandToString(srcs[i]);
        if (i != srcs.size() - 1) { srcs_str += ", "; }
    }
    if (op == OP_CALL) { srcs_str += ")"; }
    return dest + eq_str + op_str + " " + srcs_str;
}

string Block::ToString()
{
    string instrs_str = "";
    for (Instruction& instr : instrs) 
    { instrs_str += "\t" + instr.ToString() + "\n"; }
    return label + ":\n" + instrs_str;
}

string Function::ToString()
{
    string params_str = "";
    for (size_t i = 0; i < params.size(); ++i)
    {
        params_str += params[i] + ": i32";
        if (i != params.size() - 1) { params_str += ", "; } 
    }
    string type_str;
    switch (type)
    {
        case TYPE_INT: type_str = ": i32"; break;
        case TYPE_VOID: type_str = ""; break;
        default: assert(false); break;
    }
    string alloc_instrs_str = "";
    for (Instruction& instr : alloc_instrs)
    { alloc_instrs_str += "\t" + instr.ToString() + "\n"; }
    if (alloc_instrs.size() > 0) { alloc_instrs_str += "\n"; }
    string blocks_str = "";
    for (size_t i = 0; i < blocks.size(); ++i) 
    { 
        blocks_str += blocks[i].ToString(); 
        if (i != blocks.size() - 1) { blocks_str += "\n"; }
    }
    return "fun " + name + "(" + params_str + ")" + type_str 
        + "\n{\n" + alloc_instrs_str + blocks_str + "}\n";
}

string IR::ToString()
{
    string funcs_str = "";
    for (size_t i = 0; i < funcs.size(); ++i)
    {
        funcs_str += funcs[i].ToString();
        if (i != funcs.size() - 1) { funcs_str += "\n"; }
    }
    return funcs_str;
}

IR IRGenerator::Generate(ast::CompUnitAST& ast)
{
    Visit(ast);
    return move(ir);
}

void IRGenerator::Visit(ast::BaseAST& node) 
{
    if (auto* list = dynamic_cast<ast::ListAST*>(&node)) 
    { Visit(*list); } 
    else if (auto* compUnit = dynamic_cast<ast::CompUnitAST*>(&node)) 
    { Visit(*compUnit); } 
    else if (auto* funcDef = dynamic_cast<ast::FuncDefAST*>(&node)) 
    { Visit(*funcDef); } 
    else if (auto* param = dynamic_cast<ast::ParamAST*>(&node)) 
    { Visit(*param); }
    else if (auto* stmt = dynamic_cast<ast::StmtAST*>(&node)) 
    { Visit(*stmt); }
    else if (auto* assign_stmt = dynamic_cast<ast::AssignStmtAST*>(&node)) 
    { Visit(*assign_stmt); }
    else if (auto* def_stmt = dynamic_cast<ast::DefStmtAST*>(&node)) 
    { Visit(*def_stmt); }
    else if (auto* if_stmt = dynamic_cast<ast::IfStmtAST*>(&node)) 
    { Visit(*if_stmt); }
    else if (auto* while_stmt = dynamic_cast<ast::WhileStmtAST*>(&node)) 
    { Visit(*while_stmt); }
    else if (auto* expr = dynamic_cast<ast::ExprAST*>(&node)) 
    { Visit(*expr); }
    else if (auto* binary_expr = dynamic_cast<ast::BinaryExprAST*>(&node)) 
    { Visit(*binary_expr); }
    else if (auto* unary_expr = dynamic_cast<ast::UnaryExprAST*>(&node)) 
    { Visit(*unary_expr); }
    else if (auto* call_expr = dynamic_cast<ast::CallExprAST*>(&node)) 
    { Visit(*call_expr); }
    else if (auto* arg = dynamic_cast<ast::ArgAST*>(&node)) 
    { Visit(*arg); }
    else if (auto* id = dynamic_cast<ast::IdAST*>(&node)) 
    { Visit(*id); }
     else if (auto* num = dynamic_cast<ast::NumAST*>(&node)) 
    { Visit(*num); }
}

void IRGenerator::Visit(ast::ListAST& node) 
{
    for (auto& item : node.items) { Visit(*item); }
}

void IRGenerator::Visit(ast::CompUnitAST& node) 
{
    ir.funcs.clear();
    Visit(*node.func_defs);
    symbol_table.Lookup("main");
    symbol_table.Clear();
}

void IRGenerator::Visit(ast::FuncDefAST& node)
{
    local_symbol_cnt = 0;
    switch (node.type)
    {
        case ast::TYPE_INT: func.type = TYPE_INT; break;
        case ast::TYPE_VOID: func.type = TYPE_VOID; break;
        default: assert(false); break;
    }
    is_rval = false;
    Visit(*node.id);
    symbol_table.Insert(id, -1);
    func.name = GetNamedSymbol(id);
    is_in_block_stmt = true;
    symbol_table.EnterScope();
    if (node.params) { Visit(*node.params); }
    block.label = GetTempSymbol();
    if (node.stmts) { Visit(*node.stmts); }
    if (func.type == TYPE_INT) { operand1 = 0 ;}
    PushRetInstr(); 
    symbol_table.ExitScope();
    PushFunc();
}

void IRGenerator::Visit(ast::ParamAST& node)
{
    is_rval = false;
    Visit(*node.id);
    PushParam();
    operand1 = "@" + id;
    symbol_table.Insert(id, local_symbol_cnt++);
    id = GetNamedSymbol(id);
    operand2 = id;
    PushAllocInstr();
    PushStoreInstr();
}

void IRGenerator::Visit(ast::StmtAST& node)
{
    switch (node.stmt_type)
    {
        case ast::STMT_EXPR:
        case ast::STMT_ASSIGN:
        case ast::STMT_DEF:
        case ast::STMT_IF:
        case ast::STMT_WHILE:
            Visit(*node.stmt); break;
        case ast::STMT_BLOCK:
            VisitBlockStmt(node); break;
        case ast::STMT_BREAK:
            VisitBreakStmt(); break;
        case ast::STMT_CONT:
            VisitContStmt(); break;
        case ast::STMT_RET:
            VisitRetStmt(node); break;
        case ast::STMT_NULL: break;
        default: assert(false); break;
    }
}

void IRGenerator::Visit(ast::AssignStmtAST& node)
{
    Visit(*node.expr);
    is_rval = false;
    Visit(*node.id);
    id = GetNamedSymbol(id);
    operand2 = id;
    PushStoreInstr();
}

void IRGenerator::Visit(ast::DefStmtAST& node)
{
    assert(is_in_block_stmt);
    Visit(*node.expr);
    is_rval = false;
    Visit(*node.id);
    symbol_table.Insert(id, local_symbol_cnt++);
    id = GetNamedSymbol(id);
    operand2 = id;
    PushAllocInstr();
    PushStoreInstr();
}

void IRGenerator::Visit(ast::IfStmtAST& node)
{
    Visit(*node.expr);
    string then_label, else_label, end_label;
    then_label = GetTempSymbol();
    if (node.else_stmt) { else_label = GetTempSymbol(); }
    end_label = GetTempSymbol();
    operand2 = then_label;
    operand3 = node.else_stmt? else_label : end_label;
    PushBrInstr();
    block.label = then_label;
    is_in_block_stmt = false;
    Visit(*node.then_stmt);
    operand1 = end_label;
    PushJmpInstr();
    if (node.else_stmt)
    {
        block.label = else_label;
        is_in_block_stmt = false;
        Visit(*node.else_stmt);
        operand1 = end_label;
        PushJmpInstr();
    }
    block.label = end_label;
    is_in_block_stmt = true;
}

void IRGenerator::Visit(ast::WhileStmtAST& node)
{
    while_entry_labels.push_back(GetTempSymbol());
    string while_body_label = GetTempSymbol();
    while_end_labels.push_back(GetTempSymbol());
    operand1 = while_entry_labels.back();
    PushJmpInstr();
    block.label = while_entry_labels.back();
    Visit(*node.expr);
    operand2 = while_body_label;
    operand3 = while_end_labels.back();
    PushBrInstr();
    block.label = while_body_label;
    is_in_block_stmt = false;
    Visit(*node.stmt);
    operand1 = while_entry_labels.back();
    PushJmpInstr();
    block.label = while_end_labels.back();
    is_in_block_stmt = true;
    while_entry_labels.pop_back();
    while_end_labels.pop_back();
}

void IRGenerator::Visit(ast::ExprAST& node)
{
    is_rval = true;
    Visit(*node.expr);
}

void IRGenerator::Visit(ast::BinaryExprAST& node)
{
    switch (node.op)
    {
        case ast::OP_PLUS:
        case ast::OP_MINUS:
        case ast::OP_TIMES:
        case ast::OP_DIV:
        case ast::OP_MOD:
        case ast::OP_LT:
        case ast::OP_GT:
        case ast::OP_LE:
        case ast::OP_GE:
        case ast::OP_EQ:
        case ast::OP_NE:
            VisitBinaryExpr(node); break;
        case ast::OP_AND:
        case ast::OP_OR:
            VisitShortCircuitExpr(node); break;
        default:
            assert(false); break;
    }
}

void IRGenerator::Visit(ast::UnaryExprAST& node)
{
    Visit(*node.expr);
    switch (node.op)
    {
        case ast::OP_PLUS: 
             break;
        case ast::OP_MINUS: 
            op = OP_SUB;
            operand2 = move(operand1);
            operand1 = 0;
            PushBinaryInstr();
            break;
        case ast::OP_NOT: 
            op = OP_EQ;
            operand2 = 0;
            PushBinaryInstr();
            break;
        default: 
            assert(false); break;
    }
}

void IRGenerator::Visit(ast::CallExprAST& node)
{
    vector<Operand> old_args = move(call_operands);
    call_operands.clear();
    is_rval = false;
    Visit(*node.id);
    operand1 = GetNamedSymbol(id);
    PushCallOperand();
    is_rval = true;
    if (node.args) { Visit(*node.args); }
    PushCallInstr();
    call_operands = move(old_args);
}

void IRGenerator::Visit(ast::ArgAST& node)
{
    Visit(*node.expr);
    PushCallOperand();
}

void IRGenerator::Visit(ast::IdAST& node)
{
    if (is_rval)
    {
        operand1 = GetNamedSymbol(node.name);
        PushLoadInstr();
    }
    else { id = node.name; }
}

void IRGenerator::Visit(ast::NumAST& node)
{
    operand1 = node.val;
}

void IRGenerator::VisitBlockStmt(ast::StmtAST& node)
{
    is_in_block_stmt = true;
    symbol_table.EnterScope();
    Visit(*node.stmt);
    symbol_table.ExitScope();
}

void IRGenerator::VisitBreakStmt()
{
    assert(!while_end_labels.empty());
    operand1 = while_end_labels.back();
    PushJmpInstr();
    block.label = GetTempSymbol();
}

void IRGenerator::VisitContStmt()
{
    assert(!while_entry_labels.empty());
    operand1 = while_entry_labels.back();
    PushJmpInstr();
    block.label = GetTempSymbol();
}

void IRGenerator::VisitRetStmt(ast::StmtAST& node)
{
    if (func.type != TYPE_VOID) { assert(node.stmt); }
    else { assert(!node.stmt); }
    if (func.type != TYPE_VOID) { Visit(*node.stmt); }
    PushRetInstr();
    block.label = GetTempSymbol();
}

void IRGenerator::VisitBinaryExpr(ast::BinaryExprAST& node)
{
    Visit(*node.expr1);
    Operand operand = move(operand1);
    Visit(*node.expr2);
    operand2 = move(operand1);
    operand1 = move(operand);
    switch (node.op)
    {
        case ast::OP_PLUS: op = OP_ADD; break;
        case ast::OP_MINUS: op = OP_SUB; break;
        case ast::OP_TIMES: op = OP_MUL; break;
        case ast::OP_DIV: op = OP_DIV; break;
        case ast::OP_MOD: op = OP_MOD; break;
        case ast::OP_LT: op = OP_LT; break;
        case ast::OP_GT: op = OP_GT; break;
        case ast::OP_LE: op = OP_LE; break;
        case ast::OP_GE: op = OP_GE; break;
        case ast::OP_EQ: op = OP_EQ; break;
        case ast::OP_NE: op = OP_NE; break;
        default: assert(false); break;
    }
    PushBinaryInstr();
}

void IRGenerator::VisitShortCircuitExpr(ast::BinaryExprAST& node)
{
    switch (node.op)
    {
        case ast::OP_AND: operand1 = 0; break;
        case ast::OP_OR: operand1 = 1; break;
        default: assert(false); break;
    }
    string res = GetTempSymbol();
    id = res;
    operand2 = res;
    PushAllocInstr();
    PushStoreInstr();
    Visit(*node.expr1);
    string then_label = GetTempSymbol();
    string end_label = GetTempSymbol();
    switch (node.op)
    {
        case ast::OP_AND: 
            operand2 = then_label;
            operand3 = end_label; 
            break;
        case ast::OP_OR: 
            operand2 = end_label; 
            operand3 = then_label; 
            break;
        default: 
            assert(false); break;
    }
    PushBrInstr();
    block.label = then_label;
    Visit(*node.expr2);
    operand2 = res;
    PushStoreInstr();
    operand1 = end_label;
    PushJmpInstr();
    block.label = end_label;
    operand1 = res;
    PushLoadInstr();
}

void IRGenerator::PushFunc()
{
    ir.funcs.push_back(move(func));
    func.params.clear();
    func.alloc_instrs.clear();
    func.blocks.clear();
}

void IRGenerator::PushParam()
{
    func.params.push_back("@" + id);
}

void IRGenerator::PushBlock()
{
    func.blocks.push_back(move(block));
    block.instrs.clear();
}

void IRGenerator::PushInstr()
{
    block.instrs.push_back(move(instr));
    instr.srcs.clear();
}

void IRGenerator::PushOperand()
{
    instr.srcs.push_back(move(operand1));
}

void IRGenerator::PushCallOperand()
{
    call_operands.push_back(move(operand1));
}

void IRGenerator::PushBinaryInstr()
{
    instr.op = op;
    string dest = GetTempSymbol();
    instr.dest = dest;
    PushOperand();
    operand1 = move(operand2);
    PushOperand();
    PushInstr();
    operand1 = dest;
}

void IRGenerator::PushLoadInstr()
{
    instr.op = OP_LOAD;
    string dest = GetTempSymbol();
    instr.dest = dest;
    PushOperand();
    PushInstr();
    operand1 = dest;  
}

void IRGenerator::PushStoreInstr()
{
    instr.op = OP_STORE;
    instr.dest = "";
    PushOperand();
    operand1 = move(operand2);
    PushOperand();
    PushInstr();
}

void IRGenerator::PushAllocInstr()
{
    instr.op = OP_ALLOC;
    instr.dest = id;
    func.alloc_instrs.push_back(move(instr));
    instr.srcs.clear();
}

void IRGenerator::PushBrInstr()
{
    instr.op = OP_BR;
    instr.dest = "";
    PushOperand();
    operand1 = move(operand2);
    PushOperand();
    operand1 = move(operand3);
    PushOperand();
    PushInstr();
    PushBlock();
}

void IRGenerator::PushJmpInstr()
{
    instr.op = OP_JMP;
    instr.dest = "";
    PushOperand();
    PushInstr();
    PushBlock();
}

void IRGenerator::PushRetInstr()
{
    instr.op = OP_RET;
    instr.dest = "";
    if (func.type != TYPE_VOID) { PushOperand(); }
    PushInstr();
    PushBlock();
}

void IRGenerator::PushCallInstr()
{
    instr.op = OP_CALL;
    string dest = GetTempSymbol();
    instr.dest = dest;
    instr.srcs = move(call_operands);
    PushInstr();
    operand1 = dest;
}

string IRGenerator::GetTempSymbol()
{
    return "%" + to_string(local_symbol_cnt++);
}

string IRGenerator::GetNamedSymbol(string& name)
{
    string symbol = "@" + name;
    int info = symbol_table.Lookup(name);
    if (info != -1) {symbol += "_" + to_string(info); }
    return symbol;
}


